// src/main.rs
use anyhow::Result;
use libbpf_rs::{RingBufferBuilder, skel::{SkelBuilder, Skel}};
use nix::fcntl::{posix_fadvise, PosixFadviseAdvise};
use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::os::unix::io::AsRawFd;
use std::time::Duration;
use plain::Plain;

// 包含由 build.rs 生成的代码
mod prefetcher_bpf {
    include!(concat!(env!("OUT_DIR"), "/prefetcher.bpf.rs"));
}
use prefetcher_bpf::*;

// Rust 版本的事件结构体，必须与 C 版本内存布局一致
// `plain` crate 确保了这一点
#[repr(C)]
#[derive(Clone, Copy)]
struct PatternInfo {
    pub pid: u32,
    pub fault_addr: u64,
    pub stride: i64,
    pub count: u32,
}
// 安全性: 我们确保C和Rust结构体定义匹配
unsafe impl Plain for PatternInfo {}

// 进程状态机（目前很简单，后续会扩展）
struct ProcessState {
    // K: 虚拟地址起始，V: (文件句柄, 文件内偏移)
    // 这是一个简化的地址转换缓存
    vma_cache: HashMap<u64, (File, u64)>, 
}

impl ProcessState {
    fn new() -> Self {
        ProcessState {
            vma_cache: HashMap::new(),
        }
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // 1. 加载和附加 eBPF 程序
    let skel_builder = PrefetcherSkelBuilder::default();
    let open_skel = skel_builder.open()?;
    
    // 如果需要设置程序或map的参数，在这里进行
    
    let mut skel = open_skel.load()?;
    skel.attach()?;
    println!("eBPF program attached successfully.");

    // 2. 设置 Ring Buffer
    let mut ring_buf_builder = RingBufferBuilder::new(skel.maps_mut().events())?;
    
    // 存储每个被监控进程的状态
    let mut process_states: HashMap<u32, ProcessState> = HashMap::new();

    // 3. 定义事件处理回调
    let handle_event = move |data: &[u8]| {
        let mut info = PatternInfo { pid: 0, fault_addr: 0, stride: 0, count: 0 };
        if info.copy_from_bytes(data).is_err() {
            eprintln!("Failed to decode event data");
            return 0; // 返回0表示继续消费
        }

        println!(
            "Detected pattern: PID={}, Addr=0x{:x}, Stride={}, Count={}",
            info.pid, info.fault_addr, info.stride, info.count
        );
        
        let state = process_states.entry(info.pid).or_insert_with(ProcessState::new);

        // 核心预取逻辑
        if let Err(e) = prefetch(info, state) {
            eprintln!("Prefetch error for PID {}: {}", info.pid, e);
        }

        0
    };
    
    // 4. 开始轮询 Ring Buffer
    ring_buf_builder.add_callback(handle_event)?;
    let ring_buf = ring_buf_builder.build()?;
    
    println!("Waiting for events...");
    loop {
        // poll会调用回调
        if let Err(e) = ring_buf.poll(Duration::from_millis(100)) {
            eprintln!("Error polling ring buffer: {}", e);
        }
    }
}

// 预取函数
fn prefetch(info: PatternInfo, state: &mut ProcessState) -> Result<()> {
    // 预测未来N个访问地址
    const PREFETCH_PAGES: i64 = 4; // 预取多少个步长
    let page_size = 4096;

    // 1. 地址转换: 虚拟地址 -> (文件, 偏移)
    // 这是一个关键且复杂的部分。这里我们用一个简化的实现。
    // 生产级的实现需要高效地解析和缓存 /proc/[pid]/maps
    let (file, offset) = match get_file_and_offset(info.pid, info.fault_addr, &mut state.vma_cache) {
        Some(res) => res,
        None => {
            // 解析失败或地址不属于文件映射
            return Ok(());
        }
    };

    // 2. 计算预取范围并执行
    let prefetch_start_offset = (offset as i64) + info.stride;
    let prefetch_len = info.stride * PREFETCH_PAGES;

    if prefetch_len <= 0 { return Ok(()); }

    println!(
        "Prefetching for PID {}: offset={}, len={}",
        info.pid, prefetch_start_offset, prefetch_len
    );

    posix_fadvise(
        file.as_raw_fd(),
        prefetch_start_offset,
        prefetch_len,
        PosixFadviseAdvise::POSIX_FADV_WILLNEED,
    )?;

    Ok(())
}

// 简化版的地址转换函数
fn get_file_and_offset(pid: u32, vaddr: u64, cache: &mut HashMap<u64, (File, u64)>) -> Option<(File, u64)> {
    // 这是一个非常简化的实现，实际中需要解析/proc/[pid]/maps
    // 并且正确处理VMA的边界。
    // 为了演示，我们假设每次都重新解析。
    // 生产环境中必须缓存maps内容。
    let maps_path = format!("/proc/{}/maps", pid);
    let content = std::fs::read_to_string(maps_path).ok()?;
    
    for line in content.lines() {
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() < 6 { continue; }
        
        let range: Vec<&str> = parts[0].split('-').collect();
        let start = u64::from_str_radix(range[0], 16).ok()?;
        let end = u64::from_str_radix(range[1], 16).ok()?;
        
        if vaddr >= start && vaddr < end {
            let offset = u64::from_str_radix(parts[2], 16).ok()?;
            let inode = parts[4].parse::<u64>().ok()?;
            let path = parts[5];
            
            // 只关心有inode和真实路径的文件映射
            if inode == 0 || !path.starts_with('/') {
                return None;
            }
            
            let file_offset = (vaddr - start) + offset;
            let file = OpenOptions::new().read(true).open(path).ok()?;
            
            return Some((file, file_offset));
        }
    }
    None
}